package tcpserver

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/Unixeno/TheMoon/authentication"
	"github.com/Unixeno/TheMoon/common/util"
	"github.com/Unixeno/TheMoon/logging"
	"github.com/Unixeno/TheMoon/model"
	"github.com/Unixeno/TheMoon/planet/config"
	model2 "github.com/Unixeno/TheMoon/planet/orm/model"
	"github.com/Unixeno/TheMoon/planet/orm/moon"
	"github.com/Unixeno/TheMoon/pool"
	"github.com/Unixeno/TheMoon/radio/packet"
	"github.com/jinzhu/gorm"
	"io"
	"net"
	"sync"
	"time"
)

var ServerShutdownError = fmt.Errorf("server shutdown")

type TCPServer struct {
	config     *ServerConfig
	serverConn net.Listener
	recvQueue  chan *model.PlanetMessage
	clientMap  sync.Map // map[string]net.Conn
}

type ServerConfig struct {
	Addr           string        // 监听地址
	ReadDeadline   time.Duration // 读取超时时间
	WriteDeadline  time.Duration // 发送超时时间
	CAFile         string        // 用于客户端验证的ca文件，不设置时使用系统root ca
	CertFile       string        // 服务端证书文件
	PrivateKeyFile string        // 服务端证书密钥
	VerifyClient   bool          // 是否严格验证客户端，为true时会严格验证node id和客户端证书的CN字段
}

var logger = logging.GetLogger("tcp-server")

func CreateServer(config *ServerConfig) *TCPServer {
	return &TCPServer{
		config:    config,
		recvQueue: make(chan *model.PlanetMessage, 16),
	}
}

func (s *TCPServer) Init() error {
	logger.Debug("init tcp server")
	cer, err := tls.LoadX509KeyPair(s.config.CertFile, s.config.PrivateKeyFile)
	if err != nil {
		logger.WithError(err).Error("failed to load server certificate")
		return err
	} else {
		logger.WithField("cert", s.config.CertFile).WithField("private", s.config.PrivateKeyFile).
			Debug("server certificate successfully loaded")
	}

	clientAuth := tls.VerifyClientCertIfGiven
	if s.config.VerifyClient {
		// 严格验证客户端
		clientAuth = tls.RequireAndVerifyClientCert
		logger.Info("client verification enabled")
	} else {
		logger.Info("client verification disabled")
	}

	clientCAPool := x509.NewCertPool()
	if 0 != len(s.config.CAFile) {
		// 读取ca用于客户端校验
		ca, err := authentication.ReadCertificate(s.config.CAFile)
		if err != nil {
			logger.WithField("caFile", s.config.CAFile).Error("filed to load ca for client authentication")
			return err
		} else {
			logger.WithField("caFile", s.config.CAFile).Debug("load ca file succeeded")
			clientCAPool.AddCert(ca)
		}
	} else {
		// 未指定ca时使用系统ca文件
		logger.Debug("ca file not set, use system ca instead")
		clientCAPool, _ = x509.SystemCertPool() // 函数只在windows上返回错误
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cer},
		ClientCAs:    clientCAPool,
		ClientAuth:   clientAuth,
		ServerName:   config.SysConfig.ServerID,
	}

	logger.Debug("listen at ", s.config.Addr)
	conn, err := tls.Listen("tcp", s.config.Addr, tlsConfig)
	if err != nil {
		logger.WithError(err).Error()
		return err
	}
	s.serverConn = conn
	return nil
}

func (s *TCPServer) Start() {
	for {
		conn, err := s.serverConn.Accept()
		if err != nil {
			return
		} else {
			go s.handleRecv(conn)
		}
	}
}

func (s *TCPServer) handleRecv(conn net.Conn) {
	defer conn.Close()
	moonID, err := finishHandshake(conn)

	// 同一个客户端不能连两次，关闭新的连接
	if _, ok := s.clientMap.LoadOrStore(moonID, conn); ok == true {
		logger.WithField("moonID", moonID).Error("connection already exist")
		return
	}

	if err != nil {
		return
	}
	logger.WithField("MoonID", moonID).Debug("connection established with client")

	for {
		message, err := packet.ReadPacket(conn, s.config.ReadDeadline)
		if err != nil {
			if err != io.EOF {
				logger.WithError(err).WithField("client", conn.RemoteAddr()).Error("error while receiving")
			} else {
				logger.WithField("remote", conn.RemoteAddr()).Debug("connection closed by peer")
			}
			s.unregisterMoon(moonID)
			return
		}
		s.recvQueue <- &model.PlanetMessage{
			Message: message,
			MoonID:  moonID,
		}
	}
}

func (s *TCPServer) unregisterMoon(moonID string) {
	logger.WithField("moonID", moonID).Debug("unregister moon")
	s.clientMap.Delete(moonID)
}

func finishHandshake(conn net.Conn) (moonID string, err error) {
	err = nil
	if err = conn.(*tls.Conn).Handshake(); err != nil {
		logger.WithError(err).WithField("client", conn.RemoteAddr()).Error("client handshake failed")
		return
	}

	// 从证书Subject Common Name获取节点的ID
	state := conn.(*tls.Conn).ConnectionState()

	var ok bool
	moonID = state.PeerCertificates[0].Subject.CommonName
	if _, ok = util.ParseMoonID(moonID); !ok {
		logger.WithField("MoonID", moonID).Error("invalid MoonID or MoonID not provided in client certificate, shutdown connection...")
		return
	}

	dao := moon.New()
	moonDetails := model2.Moon{}
	if moonDetails, err = dao.Get(moonID); err != nil {
		if gorm.IsRecordNotFoundError(err) {
			logger.WithField("id", moonID).WithField("moon", conn.RemoteAddr()).Error("unauthorized moon id")
		} else {
			logger.WithError(err).Error("client handshake failed")
		}
		return
	} else {
		// moon第一次链接握手时需要更新节点状态
		if moonDetails.State == model2.MoonCreated {
			logger.WithField("id", moonID).Debug("first handshake, update state and signature")
			moonDetails.Signature = hex.EncodeToString(state.PeerCertificates[0].Signature)
			moonDetails.State = model2.MoonConnected
			if err = dao.Update(moonDetails); err != nil {
				logger.WithError(err).Error("update moon state failed")
				return
			}
		}
	}
	return
}

func (s *TCPServer) Recv(ctx context.Context) (*model.PlanetMessage, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case message := <-s.recvQueue:
		if message == nil {
			return nil, ServerShutdownError
		}
		return message, nil
	}
}

func (s *TCPServer) Send(ctx context.Context, message *model.PlanetMessage) error {
	logger.Debug("send to ", message.MoonID)
	conn, ok := s.clientMap.Load(message.MoonID)
	if !ok {
		return errors.New("moon not connected")
	}
	moonConn := conn.(net.Conn)

	buffer := pool.Fixed64KPool.Get().([]byte)

	defer pool.Fixed64KPool.Put(buffer)

	n, err := packet.WritePacket(buffer, message.Message)
	if err == packet.BufferTooSmall {
		logger.Error("message too long, ignore")
		return err
	}

	if _, err := moonConn.Write(buffer[:n]); err != nil {
		logger.WithError(err).WithField("moonID", message.MoonID).Error("failed to send to moon")
		return err
	}

	return nil
}
